#include "filter.h"
#include "param.h"
#include "tinyxml2.h"
#include "light_detect.h"

std::vector<ArmorAnnotation> parseXmlFile(const std::string &xml_path) {
  std::vector<ArmorAnnotation> annotations;
  tinyxml2::XMLDocument doc;
  // 读取 xml
  if (doc.LoadFile(xml_path.c_str()) != tinyxml2::XML_SUCCESS) {
    std::cerr << "读取 xml 文件失败！" << std::endl;
    return annotations;
  }

  tinyxml2::XMLElement *root = doc.FirstChildElement("annotation");
  if (!root)
    return annotations;

  // 遍历 object 节点
  for (auto obj = root->FirstChildElement("object"); obj != nullptr;
       obj = obj->NextSiblingElement("object")) {
    std::string name = obj->FirstChildElement("name")->GetText();
    if (name != "armor") // 仅提取 armor 节点
      continue;

    ArmorAnnotation ann;
    // 获取颜色，灰色跳过
    std::string color = obj->FirstChildElement("armor_color")->GetText();
    if (color == "red")
      ann.color = param::RED;
    else if (color == "blue")
      ann.color = param::BLUE;
    else
      continue;

    // 获取种类
    std::string class_text = obj->FirstChildElement("armor_class")->GetText();
    if (class_text == "none")
      ann.class_id = 0;
    else
      ann.class_id = std::stoi(class_text);

    // 获取 bbox
    auto bbox_ele = obj->FirstChildElement("bndbox");
    int xmin = std::stoi(bbox_ele->FirstChildElement("xmin")->GetText());
    int xmax = std::stoi(bbox_ele->FirstChildElement("xmax")->GetText());
    int ymin = std::stoi(bbox_ele->FirstChildElement("ymin")->GetText());
    int ymax = std::stoi(bbox_ele->FirstChildElement("ymax")->GetText());
    ann.bbox = cv::Rect(xmin, ymin, xmax - xmin, ymax - ymin);

    annotations.push_back(ann);
  }

  return annotations;
}

bool findAndWarpArmor(cv::Mat &src, cv::Mat &dst, bool color) {
  ArmorDetect detector(color);
  cv::Mat result;
  detector.detect(src, result);
  if (result.empty())
    return false;
  
  dst = result;
  return true;
}
